#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import tempfile
import re
import logging
import uuid
import io
import shutil
from typing import Dict, List, Any, BinaryIO, Tuple, Optional
from pathlib import Path

from document_processor.engines.base_engine import BaseEngine
from markitdown import MarkItDown

logger = logging.getLogger(__name__)

class MarkItDownEngine(BaseEngine):
    """使用MarkItDown工具转换文档为Markdown的引擎"""
    
    def __init__(self, config: Dict[str, Any]):
        """初始化MarkItDown引擎
        
        Args:
            config: 配置字典
        """
        super().__init__(config)
        
        # 初始化转换器
        try:
            self.converter = MarkItDown(
                enable_plugins=True,  # 启用插件支持
                enable_builtins=True  # 启用内置转换器
            )
            logger.info("已成功初始化MarkItDown转换器")
        except Exception as e:
            logger.error(f"初始化MarkItDown转换器失败: {str(e)}")
            raise

    async def convert_to_markdown(
        self, 
        file_content: BinaryIO, 
        file_name: str,
        image_option: str = "remove",
        work_dir: Optional[str] = None,
        md_filename: Optional[str] = None
    ) -> Dict[str, Any]:
        """使用MarkItDown将文档转换为Markdown格式
        
        Args:
            file_content: 文件内容（二进制流）
            file_name: 文件名
            image_option: 图片处理选项，只支持 "remove"（移除图片）
            work_dir: 工作目录路径，如果提供则使用此目录，否则创建临时目录
            
        Returns:
            包含以下键的字典:
                - text: 提取的纯文本
                - markdown: 转换后的Markdown文本
                - images: 图片信息列表
                - temp_dir: 临时目录路径
        """
        # 检查图片处理选项
        if image_option != "remove":
            logger.warning("MarkItDown引擎只支持移除图片选项")
            image_option = "remove"
            
        # 检查文件格式
        file_ext = os.path.splitext(file_name)[1].lower()
        supported_formats = self.config["document"]["supported_formats"]
        if file_ext not in supported_formats:
            raise ValueError(f"不支持的文件格式: {file_ext}，支持的格式: {', '.join(supported_formats)}")
            
        # 使用提供的工作目录或创建临时目录
        temp_dir = work_dir if work_dir else tempfile.mkdtemp()
        if not os.path.exists(temp_dir):
            os.makedirs(temp_dir)
            
        try:
            # 保存文件到临时目录
            temp_file = os.path.join(temp_dir, file_name)
            with open(temp_file, 'wb') as f:
                f.write(file_content.read())
            
            logger.info(f"开始处理文档: {file_name}")
            
            try:
                # 使用MarkItDown转换文档
                result = self.converter.convert(temp_file)
                markdown_text = result.text_content
                logger.info("文档转换成功")
                
                # 清理可能存在的base64图片引用
                markdown_text = re.sub(r'!\[.*?\]\(data:image/[^)]+\)', '', markdown_text)
                
                # 保存Markdown内容到临时文件
                output_file = os.path.join(temp_dir, f"{Path(file_name).stem}.md")
                with open(output_file, 'w', encoding='utf-8') as f:
                    f.write(markdown_text)
                
                return {
                    "text": re.sub(r'!\[.*?\]\(.*?\)', '', markdown_text),  # 纯文本（移除图片引用）
                    "markdown": markdown_text,
                    "images": [],  # 不处理图片
                    "temp_dir": temp_dir
                }
                
            except Exception as e:
                logger.error(f"转换失败: {str(e)}")
                raise
                
        except Exception as e:
            logger.error(f"处理文件出错: {str(e)}")
            # 如果是临时目录，则在出错时清理
            if not work_dir and temp_dir and os.path.exists(temp_dir):
                shutil.rmtree(temp_dir, ignore_errors=True)
            raise

    async def extract_images(
        self, 
        file_content: BinaryIO,
        file_name: str
    ) -> Tuple[List[Dict[str, Any]], Optional[str]]:
        """从文档中提取图片
        
        Args:
            file_content: 文件内容（二进制流）
            file_name: 文件名
            
        Returns:
            Tuple[图片信息列表, 临时目录路径]
        """
        # 创建临时目录
        temp_dir = tempfile.mkdtemp()
        return [], temp_dir

    async def handle_images(
        self, 
        images: List[Dict[str, Any]], 
        markdown_text: str,
        uploader
    ) -> str:
        """处理图片（上传到OBS并更新Markdown中的图片引用）
        
        Args:
            images: 图片信息列表
            markdown_text: Markdown文本
            uploader: OBS上传器实例
            
        Returns:
            更新后的Markdown文本（图片引用已替换为OBS URL）
        """
        # 直接返回原始文本，不处理图片
        return markdown_text 